package model

import (
	"common/comCopier"
	"common/dao/model"
	"common/utils"
	"context"
	"encoding/json"
	"fmt"
	"gateway/chatUtils"
	"github.com/go-deepseek/deepseek"
	"github.com/go-deepseek/deepseek/request"
	"gorm.io/gorm"
	"io"
	"net/http"
	"time"

	"gateway/internal/svc"
	"gateway/internal/types"

	"github.com/zeromicro/go-zero/core/logx"
)

type ChatStreamLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
	w      http.ResponseWriter
	r      *http.Request
}

func NewChatStreamLogic(ctx context.Context, svcCtx *svc.ServiceContext, w http.ResponseWriter, r *http.Request) *ChatStreamLogic {
	return &ChatStreamLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
		w:      w,
		r:      r,
	}
}

func (l *ChatStreamLogic) ChatStream(req *types.ChatReq) error {
	var (
		modelData           = model.Model{}
		historyMessage      []*request.Message      // 历史会话
		conversations       []model.Conversation    // 数据库历史会话
		ConversationGroup   model.ConversationGroup //历史会话组
		nowUserConversation model.Conversation      //当前的会话返回
		nowConversation     model.Conversation      //当前的会话返回
		user                = model.User{}
		session             chatUtils.SessionData
	)

	tokenInfo := l.ctx.Value("token")
	err := comCopier.CopyWithOption(&user, tokenInfo)
	if err != nil {
		return err
	}
	l.w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	l.w.Header().Set("Transfer-Encoding", "chunked")
	session.SessionId = utils.GetUuid()
	session.HttpCtx = l.r.Context()
	session.IsSession = true
	chatUtils.SetDataManager(&session)
	fmt.Println(session.SessionId)
	sessionByte, err := json.Marshal(session)
	if err != nil {
		return err
	}
	chatUtils.SendStreamEvent(l.w, l.r, string(sessionByte))
	if req.ConversationGroupId > 0 {
		err := l.svcCtx.DB.Table(model.TableNameConversation).Transaction(func(tx *gorm.DB) error {
			return tx.Where("conversationId = ?", req.ConversationGroupId).Find(&conversations).Error
		})
		if err != nil {
			return err
		}
		err = l.svcCtx.DB.Table(model.TableNameConversationGroup).Transaction(func(tx *gorm.DB) error {
			return tx.Where("id = ?", req.ConversationGroupId).Find(&ConversationGroup).Error
		})

		if err != nil {
			return err
		}

	} else {
		ConversationGroup.UserID = user.ID
		ConversationGroup.CreateTime = time.Now().Unix()
		ConversationGroup.UpdateTime = time.Now().Unix()
		err := l.svcCtx.DB.Table(model.TableNameConversationGroup).Transaction(func(tx *gorm.DB) error {
			return tx.Create(&ConversationGroup).Error
		})
		if err != nil {
			return err
		}

	}

	if len(conversations) > 0 {
		for _, conversation := range conversations {
			historyMessage = append(historyMessage, &request.Message{
				Role:    conversation.Role,
				Content: conversation.ConversationContent,
			})
		}
	}

	historyMessage = append(historyMessage, &request.Message{
		Role:    "user",
		Content: req.UserMessage,
	})
	chatReq := &request.ChatCompletionsRequest{
		Model:    req.Model,
		Stream:   true,
		Messages: historyMessage,
	}
	err = l.svcCtx.DB.Table(model.TableNameModel).Transaction(func(tx *gorm.DB) error {
		return tx.Where("id = ?", req.ModelId).Find(&modelData).Error
	})
	if err != nil {
		return err
	}
	cli, _ := deepseek.NewClient(modelData.APIKey)
	ctx := context.Background()
	session.SessionCtx = ctx
	reasoner, err := cli.StreamChatCompletionsReasoner(ctx, chatReq)
	if err != nil {
		return err
	}

	startThinkTime := time.Now().Unix()
	for {
		data := chatUtils.GetDataManager(session.SessionId)
		if data.IsSession == false {
			chatUtils.SetDataManager(data)
			break
		}
		chatResp, err := reasoner.Read()
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		if chatResp.Choices[0].Delta.ReasoningContent != "" {
			nowConversation.ReasoningContent += chatResp.Choices[0].Delta.ReasoningContent

		} else if chatResp.Usage != nil && chatResp.Usage.TotalTokens > 0 {
			nowConversation.Tokens = int32(chatResp.Usage.TotalTokens)
		} else {
			if nowConversation.ThinkingTime == 0 {
				nowConversation.ThinkingTime = int32(time.Now().Unix() - startThinkTime)
			}
			nowConversation.ConversationContent += chatResp.Choices[0].Delta.Content
		}
		marshal, err := json.Marshal(nowConversation)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		chatUtils.SendStreamEvent(l.w, l.r, string(marshal))
	}
	if req.ConversationGroupId == 0 {
		marshal, err := json.Marshal(ConversationGroup)
		if err != nil {
			return err
		}
		chatUtils.SendStreamEvent(l.w, l.r, string(marshal))
	}
	chatUtils.SendStreamEvent(l.w, l.r, "[DONE]")
	nowUserConversation.CreateTime = time.Now().Unix()
	nowUserConversation.UpdateTime = time.Now().Unix()
	nowUserConversation.UserID = user.ID
	if len(conversations) == 0 {

		nowUserConversation.ParentID = 0
	} else {
		nowUserConversation.ParentID = conversations[len(conversations)-1].ID
	}

	nowUserConversation.ConversationID = ConversationGroup.ID
	nowUserConversation.Role = "user"
	nowUserConversation.ConversationContent = req.UserMessage
	err = l.svcCtx.DB.Table(model.TableNameConversation).Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(&nowUserConversation).Error; err != nil {
			return err
		}
		nowConversation.ParentID = nowUserConversation.ID
		nowConversation.UserID = user.ID
		nowConversation.ConversationID = ConversationGroup.ID
		nowConversation.Role = "assistant"
		nowConversation.CreateTime = time.Now().Unix()
		nowConversation.UpdateTime = time.Now().Unix()
		if err := tx.Create(&nowConversation).Error; err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return err
	}
	session.IsSession = false
	chatUtils.SetDataManager(&session)
	return nil
}
